import torch
import pickle
import gzip
from torch.utils.data import DataLoader,TensorDataset

def load_data():
    f = open('../data/mnist.pkl', 'rb')
    training_data, validation_data, test_data = pickle.load(f, encoding='bytes')  # 按字节排
    f.close()
    return (training_data, validation_data, test_data)

def load_data_wrapper():

    tr_d, va_d, te_d = load_data()

    #把数据转换成tensor张量
    train_images, train_labels = torch.tensor(tr_d[0], dtype=torch.float32), torch.tensor(tr_d[1], dtype=torch.long)
    val_images, val_labels = torch.tensor(va_d[0], dtype=torch.float32), torch.tensor(va_d[1], dtype=torch.long)
    test_images, test_labels = torch.tensor(te_d[0], dtype=torch.float32), torch.tensor(te_d[1], dtype=torch.long)

    # 创建数据集对象
    train_dataset = TensorDataset(train_images, train_labels)
    val_dataset = TensorDataset(val_images, val_labels)
    test_dataset = TensorDataset(test_images, test_labels)

    # 创建DataLoader对象
    train_loader = DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)
    val_loader = DataLoader(dataset=val_dataset, batch_size=16, shuffle=False)
    test_loader = DataLoader(dataset=test_dataset, batch_size=16, shuffle=False)

    return (train_loader, val_loader, test_loader)

